"""
概率问题 - 贝叶斯公式
// 有条件概率公式: 条件概率描述在已知事件 B 发生的情况下，事件 A 发生的概率。
P(A|B) = P(A ∩ B) / P(B)   ----P(B) > 0
// 全概率公式: 全概率公式用于计算一个事件 A 的总概率。
p(A) = ∑n/i=1 P(A|Bi) * P(Bi)  ----n指的是划分事件的总数
// 贝叶斯公式: 贝叶斯公式是条件概率和全概率公式的结合，常用于求解“逆概率”（即已知结果求原因的概率）
- 全概率为分母
- 分子根据事件识别 各条件概率相乘
P(Bi|A) = [P(A|Bi) * P(Bi)] / P(A)
"""


class BayesianCalculator:
    def __init__(self):
        self.groups = {}
        self.evidence = None
        self.total_probability = None

    def add_group(self, name: str, prior: float, likelihoods: dict):
        """
        添加一个可能的原因组

        :param name: 组名
        :param prior: 先验概率
        :param likelihoods: 证据的概率分布
        :return:
        """
        self.groups[name] = {
            "prior": prior,
            "likelihoods": likelihoods
        }

    def set_evidence(self, evidence: str):
        """
        设置观察到的证据

        :param evidence: 证据
        :return:
        """
        self.evidence = evidence
        self._calculate_total_probability()

    def _calculate_total_probability(self):
        """
        计算全概率P(证据)
        :return:
        """
        if not self.evidence:
            raise ValueError("please set the evidence first！")
        total = 0.0
        for group, data in self.groups.items():
            prior = data["prior"]
            likelihoods = data["likelihoods"].get(self.evidence, 0)
            total += prior * likelihoods
        self.total_probability = total

    def posterior(self, group: str) -> float:
        """
        计算后验概率 P(组|证据)
        :param group:
        :return:
        """
        if not self.evidence:
            raise ValueError("please set the evidence first！")
        if group not in self.groups:
            raise ValueError(f"{group}group not found！")
        data = self.groups[group]
        prior = data["prior"]
        likelihoods = data["likelihoods"].get(self.evidence, 0)

        return (prior * likelihoods) / self.total_probability

    def __str__(self):
        if not self.evidence:
            return "no evidence set！"
        output = [
            f"bayes calculator - evidence: {self.evidence}",
            f"total probability is P({self.evidence})={self.total_probability}",
            "\nposterior in each group:"
        ]
        for group in self.groups:
            post = self.posterior(group)
            output.append(f"- P({group}|{self.evidence}) = {post:.6f} ({post * 100:.2f}%)")
        return "\n".join(output)


"""
已知甲袋中有6只红球，4只白球，
乙袋中有8只红球，6只白球，
随机取一只袋，再从袋中任取一球，发现是红球，则此球来自甲袋的概率为
"""
instance = BayesianCalculator()
instance.add_group(name="甲袋", prior=0.5, likelihoods={
    "红球": 6 / 10,
    "白球": 4 / 10
})
instance.add_group(name="乙袋", prior=0.5, likelihoods={
    "红球": 8 / 14,
    "白球": 6 / 14
})
instance.set_evidence("红球")
res = instance.posterior("甲袋")
print(instance)
print(f"\nthe final calculated result is {res:.6f}")


